package learning.crf.training;

import java.text.DecimalFormat;
import java.text.NumberFormat;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

import learning.crf.training.CRF.DecodingResult;
import learning.data.Dataset;
import learning.data.document.SequenceDocument;
import learning.util.SparseVector;

public class FormatHelper {

	public static String formatDocument(SequenceDocument doc) {

        int[] columnWidths = new int[doc.tokens.length];
        for (int i = 0; i < doc.tokens.length; i++) columnWidths[i] =
            doc.tokens[i].length() + 1;

        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < columnWidths.length; i++)
            sb.append(String.format("%1$-" + (columnWidths[i]) + "s", doc.tokens[i]));
        sb.append("\n");
        for (int i = 0; i < columnWidths.length; i++)
            sb.append(String.format("%1$-" + (columnWidths[i]) + "s",
                labelText(doc.labels[i], doc.tokens[i])));
        sb.append("\n");
        
        return (sb.toString());
	}
	
	// print an example document with its predicted and expected tag sequences
	public static String formatExample(SequenceDocument doc, DecodingResult predictedParse, DecodingResult trueParse) {
		
        int[] columnWidths = new int[doc.tokens.length];
        for (int i = 0; i < doc.tokens.length; i++) columnWidths[i] =
            doc.tokens[i].length() + 1;

        StringBuilder sb = new StringBuilder();
        sb.append(  "Predicted:");
        for (int i = 0; i < columnWidths.length; i++)
            sb.append(String.format("%1$-" + (columnWidths[i]) + "s",
                labelText(predictedParse.labels[i], doc.tokens[i])));
        sb.append("\nExample:  ");
        for (int i = 0; i < columnWidths.length; i++)
            sb.append(String.format("%1$-" + (columnWidths[i]) + "s", doc.tokens[i]));
        sb.append("\nExpected: ");
        for (int i = 0; i < columnWidths.length; i++)
            sb.append(String.format("%1$-" + (columnWidths[i]) + "s",
                labelText(trueParse.labels[i], doc.tokens[i])));
        sb.append("\n");
        
        return (sb.toString());
	}
	
	private static String labelText(int label, String token) {
		//if (label == 0) return " ";
		String l = "";
		for (int i=0; i < token.length(); i++) l += "" + label; //"*";
		return l;
	}

	// print the features with highest weights for each state transition
	public static String formatParams(CRFParameters p, Dataset<SequenceDocument> dataset) {
		NumberFormat format = new DecimalFormat("0.00");
		
		String[] names = new String[p.model.numStates * p.model.numStates];
		List[] lists = new List[p.model.numStates * p.model.numStates];
		int[] widths = new int[p.model.numStates * p.model.numStates];
		
		for (int i=0; i < p.transitionParameters.length; i++) {
			for (int j=0; j < p.transitionParameters.length; j++) {
				int pos = i*p.model.numStates + j; 
				names[pos] = i + " -> " + j;
				SparseVector v = p.transitionParameters[i][j];
				List<Pair> pairs = new ArrayList<Pair>(v.num);
				
				for (int k=0; k < v.num; k++) pairs.add 
					(new Pair(v.ids[k], dataset.getFeatureName(v.ids[k]), v.vals[k]));
				
				Collections.sort(pairs);
				int w = 1;
				for (int k=0; k < Math.min(10, pairs.size()); k++)
					w = Math.max(w, pairs.get(k).feature.length()+1);
				widths[pos] = w + 6;
				lists[pos] = pairs;
			}
		}
		
		StringBuilder sb = new StringBuilder();
		for (int i=0; i < names.length; i++) {			
			sb.append(String.format("%1$-" + (widths[i]) + "s", names[i]));
		}
		sb.append("\n");
		for (int l=0; l < 10; l++) {
			for (int i=0; i < lists.length; i++) {
				if (lists[i].size() > l) {
					Pair pa = (Pair)lists[i].get(l);
					
					sb.append(String.format("%1$-" + (widths[i]) + "s",
							format.format(pa.score) + " " + pa.feature));					
				}				
			}
			sb.append("\n");
		}
		return (sb.toString());
	}
	
	// print the features with highest weights
	public static String formatVector(SparseVector v, Dataset<SequenceDocument> dataset) {
		NumberFormat format = new DecimalFormat("0.00");
		
		List<Pair> pairs = new ArrayList<Pair>(v.num);
		
		for (int k=0; k < v.num; k++) pairs.add 
			(new Pair(v.ids[k], dataset.getFeatureName(v.ids[k]), v.vals[k]));
		
		Collections.sort(pairs);

		StringBuilder sb = new StringBuilder();
		for (int i=0; i < pairs.size(); i++) {
			Pair pa = (Pair)pairs.get(i);
			
			sb.append(format.format(pa.score) + " " + pa.feature + "\n");					
		}
		return (sb.toString());
	}
	
	private static class Pair implements Comparable<Pair> {
		int id;
		String feature;
		float score;
		Pair(int id, String feature, float score) { 
			this.id = id; this.feature = feature; this.score = score; 
		}
		
		public int compareTo(Pair p) {
			if (Math.abs(score) < Math.abs(p.score)) return 1; 
			else if (Math.abs(score) > Math.abs(p.score)) return -1; 
			return 0;
		}
	}

	// print confusion matrix
	public static String formatConfusionMatrix(Evaluator evaluator) {
		DecimalFormat format = new DecimalFormat("0.00000");
		String ansiRed    = "\u0027[31m";
		String ansiNormal = "\u0027[0m";
		boolean useAnsi = false;
		
		int numLabels = evaluator.getNumLabels();
		int[][] confusionMatrix = evaluator.getConfusionMatrix();

		int[] trueCounts = new int[numLabels]; 
		for (int i=0; i < numLabels; i++)
			for (int j=0; j < numLabels; j++)
				trueCounts[i] += confusionMatrix[i][j];
		
		StringBuilder sb = new StringBuilder();
		for (int i=0; i < numLabels; i++) {
			sb.append("\t" + i);
		}
		sb.append("\n");
		
		for (int i=0; i < numLabels; i++) {
			sb.append(i);
			
			for (int j=0; j < numLabels; j++) {
				sb.append("\t");
				float f = (float)(confusionMatrix[i][j] / (double)trueCounts[i]);
				String n = format.format(f);
				if (f > 0 && useAnsi) n = ansiRed + n + ansiNormal;
				sb.append(n);
			}
			sb.append("\n");
		}
		return sb.toString();
	}
}
